from pyflowchart import Flowchart

code = """
x1 = self.first_linear(x1)
x2 = self.second_linear(x2)
for x in range(10):
    if x % 2 == 0:
        for a in range(x):
            print(a)
x = torch.cat((x1,x2),dim = -1)
ret = self.linear_stack(x)
return ret
"""
fc = Flowchart.from_code(code)
print(fc.flowchart())
